# knowledge_graph.py
import os
import threading
import json
import re
from neo4j import GraphDatabase
from openai import OpenAI
from typing import List, Tuple, Dict
from concurrent.futures import ThreadPoolExecutor, as_completed
from langchain.text_splitter import RecursiveCharacterTextSplitter
from config import API_KEY, BASE_URL, MODEL_NAME, NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD, DATABASE_NAME


class KnowledgeGraphProcessor:
    """管理知识图谱构建和查询的类。"""

    def __init__(self, uri: str = NEO4J_URI, user: str = NEO4J_USER, password: str = NEO4J_PASSWORD,
                 database: str = DATABASE_NAME, kg_cache_path: str = None):
        self.driver = None
        self.uri = uri
        self.user = user
        self.password = password
        self.database = database
        self.kg_cache_path = kg_cache_path
        self.client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
        self._connect_to_db()
        self.create_schema()

    def _connect_to_db(self):
        """连接Neo4j数据库。"""
        try:
            self.driver = GraphDatabase.driver(self.uri, auth=(self.user, self.password))
            self.driver.verify_connectivity()
            print("Successfully connected to Neo4j database.")
        except Exception as e:
            print(f"Failed to connect to Neo4j: {e}")
            self.driver = None

    def close(self):
        """关闭数据库连接。"""
        if self.driver:
            self.driver.close()
            print("Neo4j database connection closed.")

    def create_schema(self):
        """创建知识图谱模式。"""
        if not self.driver: return
        with self.driver.session(database=self.database) as session:
            try:
                session.run("CREATE CONSTRAINT FOR (n:Chunk) REQUIRE n.chunk_id IS UNIQUE")
                session.run("CREATE CONSTRAINT FOR (n:Entity) REQUIRE n.name IS UNIQUE")
                session.run("CREATE INDEX FOR (r:Relation) ON (r.name)")
                session.run("CREATE INDEX FOR (e:Entity) ON (e.name)")
                print("Neo4j schema created successfully.")
            except Exception as e:
                # 捕获已存在约束的特定错误
                if "EquivalentSchemaRuleAlreadyExists" in str(e):
                    print("Neo4j schema already exists, skipping creation.")
                else:
                    raise e

    def clear_graph(self):
        """清空知识图谱。"""
        if not self.driver: return
        with self.driver.session(database=self.database) as session:
            session.run("MATCH (n) DETACH DELETE n")
            print("Knowledge graph cleared.")

    def _extract_triplets(self, text: str) -> List[Tuple[str, str, str]]:
        """使用大语言模型从文本中提取知识三元组。"""
        prompt = f"""你是一个知识图谱提取专家。你的任务是从以下文本中提取知识三元组，格式为(实体1, 关系, 实体2, 来源)。
        请尽可能提取所有有意义的三元组。来源应该是提取三元组的原始句子。
        如果一个实体在文本中没有直接对应的另一个实体来形成三元组，则不要提取。
        文本: "{text}"
        三元组: """
        try:
            completion = self.client.chat.completions.create(
                model=MODEL_NAME,
                messages=[
                    {"role": "user", "content": prompt}
                ],
                max_tokens=2048,
                temperature=0.3
            )
            response_text = completion.choices[0].message.content.strip()
            # 解析三元组
            triplets = []
            pattern = re.compile(r"\(('.*?'),\s*('.*?'),\s*('.*?'),\s*('.*?')\)")
            matches = pattern.findall(response_text)
            for match in matches:
                triplets.append(tuple(s.strip("'") for s in match))
            return triplets
        except Exception as e:
            print(f"Failed to extract triplets: {e}")
            return []

    def load_graph_from_cache(self, current_files_metadata: List[Dict]) -> bool:
        """
        从缓存文件加载知识图谱，并验证文件元数据是否匹配。
        如果元数据不匹配，则缓存失效。
        """
        if os.path.exists(self.kg_cache_path):
            try:
                with open(self.kg_cache_path, 'r', encoding='utf-8') as f:
                    cache_data = json.load(f)

                # 检查缓存中是否存在元数据，并与当前元数据进行比对
                cached_metadata = cache_data.get('metadata')
                if cached_metadata == current_files_metadata:
                    print("发现本地知识图谱缓存，且文件未更改，正在加载...")
                    triplets = cache_data['triplets']

                    # 将缓存的三元组导入Neo4j
                    self.clear_graph()
                    if not self.driver:
                        return False
                    with self.driver.session(database=self.database) as session:
                        for s, p, o, source in triplets:
                            session.run(
                                "MERGE (s:Entity {name: $s}) "
                                "MERGE (o:Entity {name: $o}) "
                                "MERGE (s)-[:`" + p + "`]->(o)",
                                s=s, o=o
                            )
                    print("知识图谱已成功从缓存加载。")
                    return True
                else:
                    print("发现本地知识图谱缓存，但源文件已更改，将重新构建图谱。")
            except Exception as e:
                print(f"加载知识图谱缓存失败：{e}，将重新构建。")
        else:
            print("未找到本地知识图谱缓存，将进行在线构建。")

        return False

    def populate_graph(self, documents: List, event: threading.Event, files_metadata: List[Dict]) -> None:
        """异步构建知识图谱，并设置事件通知。"""
        if not self.driver:
            event.set()
            return

        print("知识图谱正在后台构建...")
        triplets_list = []
        with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
            futures = [executor.submit(self._extract_triplets, doc.page_content) for doc in documents]
            for future in as_completed(futures):
                triplets_list.extend(future.result())

        # 将三元组和文件元数据一起缓存到文件中
        cache_data = {
            "triplets": triplets_list,
            "metadata": files_metadata
        }

        parent_dir = os.path.dirname(self.kg_cache_path)
        os.makedirs(parent_dir, exist_ok=True)
        with open(self.kg_cache_path, 'w', encoding='utf-8') as f:
            json.dump(cache_data, f, ensure_ascii=False, indent=2)

        print(f"知识图谱数据已保存到缓存文件：{self.kg_cache_path}")

        # 将三元组导入Neo4j
        self.clear_graph()
        with self.driver.session(database=self.database) as session:
            for s, p, o, source in triplets_list:
                session.run(
                    "MERGE (s:Entity {name: $s}) "
                    "MERGE (o:Entity {name: $o}) "
                    "MERGE (s)-[:`" + p + "`]->(o)",
                    s=s, o=o
                )
        print("知识图谱构建完成！")
        event.set()

    def query_graph(self, query_keywords: List[str]) -> str:
        """根据关键词查询知识图谱，并返回相关信息。"""
        if not self.driver: return "Neo4j数据库未连接。"
        cypher_query = """
        MATCH (n:Entity)-[r]->(m:Entity)
        WHERE n.name IN $keywords OR m.name IN $keywords
        RETURN n.name AS start_node, type(r) AS relation, m.name AS end_node
        LIMIT 10
        """
        result_text = "相关知识点:\n"
        with self.driver.session(database=self.database) as session:
            result = session.run(cypher_query, keywords=query_keywords)
            for record in result:
                result_text += f"- {record['start_node']} -> {record['relation']} -> {record['end_node']}\n"
        return result_text